from docx import Document
import requests
import os
import base64  # 添加缺失的导入

def mermaid_to_word(mermaid_code, output_docx):
    try:
        # 1. 将 Mermaid 代码转换为 base64
        graphbytes = mermaid_code.encode("utf8")
        base64_bytes = base64.b64encode(graphbytes)
        base64_str = base64_bytes.decode("utf8")
        
        # 2. 使用 Mermaid 在线服务生成图片
        url = f"https://mermaid.ink/img/{base64_str}?type=png"
        
        # 3. 下载图片
        response = requests.get(url)
        response.raise_for_status()  # 检查请求是否成功
        
        img_path = "temp_flowchart.png"
        with open(img_path, "wb") as f:
            f.write(response.content)
        
        # 4. 插入到 Word
        doc = Document()
        doc.add_heading('流程图', level=1)
        doc.add_picture(img_path, width=docx.shared.Inches(6))  # 控制宽度
        doc.save(output_docx)
        
        print(f"成功生成 Word 文档: {output_docx}")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")
    finally:
        # 5. 清理临时文件
        if os.path.exists(img_path):
            os.remove(img_path)

# 示例调用
if __name__ == "__main__":
    mermaid_code = """
    graph TD
        A[开始] --> B{条件}
        B -->|是| C[执行操作]
        B -->|否| D[结束]
    """
    mermaid_to_word(mermaid_code, "output.docx")
